import torch
from numpy.typing import ArrayLike

from .utils import torch_distributed_get_info


class ProgressRecorder:
    def __init__(self, value: torch.Tensor | ArrayLike = 0):
        _, self.processes_num = torch_distributed_get_info()
        self._value = torch.asarray(value, dtype=torch.int64).detach().sum()

    def get(self) -> int:
        return self._value.cpu().item()

    def increase(self, value: torch.Tensor | ArrayLike = 0):
        value = torch.asarray(value, dtype=torch.int64)
        value = value.detach().sum()
        if self.processes_num > 1:
            torch.distributed.all_reduce(value, op=torch.distributed.ReduceOp.SUM)
        value = torch.asarray(value, dtype=torch.int64)
        self._value = self._value.to(value.device)
        self._value += value

    def assign(self, value: torch.Tensor | ArrayLike = 0):
        value = torch.asarray(value, dtype=torch.int64)
        self._value = value.detach().sum()
